import torch
import glob
import os
import cv2
from joblib import dump, load
import numpy as np
from sklearn.svm import SVC
from sklearn.multiclass import OneVsRestClassifier
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch import nn
from torch import optim
from PIL import Image
from torchvision import transforms
from fatigue_algo.fatigue_detector import mobilenetv3
import os

'''基于mobileNet模型的皱眉推理'''
def frown_infer_CE(image,model):
    '''
    :param image: ndarray
    :return:
    '''
    #将要检测的图片放大至224 x 224
    image = cv2.resize(image,(224,224),interpolation=cv2.INTER_CUBIC)
    w, h, c = image.shape
    image = np.resize(image, (c, w, h))
    img_tensor = torch.tensor(image, dtype=torch.float32)
    img_tensor = img_tensor.unsqueeze(0)

    # device = torch.device("cpu")
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    img_tensor = img_tensor.to(device)
    model.to(device)
    model.eval()
    out = model(img_tensor)
    pred, pred_index = out.max(axis=1)
    return pred_index
